package com.chenbingkang.bigdata.homework5;

import org.apache.commons.lang3.tuple.MutablePair;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.FlatMapFunction;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.Function2;
import scala.Tuple2;

import java.util.*;

public class Application {
    public static void main(String... args) {
        SparkConf sparkConf = new SparkConf().setMaster("local").setAppName("hello world");
        JavaSparkContext sparkContext = new JavaSparkContext(sparkConf);
        JavaRDD<String> textFile = sparkContext.textFile("text");
        List<Tuple2<String, Iterable<Pair<String, List<Pair<String, Integer>>>>>> result = textFile.map(new Function<String, Pair<String, String>>() {
            public Pair<String, String> call(String v1) throws Exception {
                String[] arr = v1.split("\\.");
                String docId = arr[0];
                String docContent = arr[1];
                return new MutablePair(docId, docContent);
            }
        }).map(new Function<Pair<String, String>, Pair<String, List<String>>>() {
            public Pair<String, List<String>> call(Pair<String, String> v1) throws Exception {
                return new MutablePair<String, List<String>>(v1.getLeft(), Arrays.asList(v1.getRight().split("\\s+")));
            }
        }).map(new Function<Pair<String, List<String>>, Pair<String, Map<String, Integer>>>() {
            public Pair<String, Map<String, Integer>> call(Pair<String, List<String>> v1) throws Exception {
                List<String> tokens = v1.getRight();
                Map<String, Integer> tokenCountMap = new HashMap<String, Integer>();
                for (String token : tokens) {
                    if (tokenCountMap.containsKey(token)) {
                        tokenCountMap.put(token, tokenCountMap.get(token) + 1);
                    } else {
                        tokenCountMap.put(token, 1);
                    }
                }
                return new MutablePair<String, Map<String, Integer>>(v1.getLeft(), tokenCountMap);

            }
        }).flatMap(new FlatMapFunction<Pair<String, Map<String, Integer>>, Pair<String, List<Pair<String, Integer>>>>() {
            public Iterator<Pair<String, List<Pair<String, Integer>>>> call(Pair<String, Map<String, Integer>> v1) throws Exception {
                String docId = v1.getLeft();
                Map<String, Integer> tokenCountMap = v1.getRight();
                List<Pair<String, List<Pair<String, Integer>>>> tokens = new ArrayList<Pair<String, List<Pair<String, Integer>>>>();
                for (Map.Entry<String, Integer> entry : tokenCountMap.entrySet()) {
                    List<Pair<String, Integer>> list = new ArrayList<Pair<String, Integer>>();
                    list.add(new MutablePair<String, Integer>(docId, entry.getValue()));
                    tokens.add(new MutablePair<String, List<Pair<String, Integer>>>(entry.getKey(), list));
                }
                return tokens.iterator();
            }
        }).groupBy(new Function<Pair<String, List<Pair<String, Integer>>>, String>() {
            public String call(Pair<String, List<Pair<String, Integer>>> v1) throws Exception {
                return v1.getKey();
            }
        }).reduceByKey(new Function2<Iterable<Pair<String, List<Pair<String, Integer>>>>, Iterable<Pair<String, List<Pair<String, Integer>>>>, Iterable<Pair<String, List<Pair<String, Integer>>>>>() {
            public Iterable<Pair<String, List<Pair<String, Integer>>>> call(Iterable<Pair<String, List<Pair<String, Integer>>>> v1, Iterable<Pair<String, List<Pair<String, Integer>>>> v2) throws Exception {
                MutablePair<String, List<Pair<String, Integer>>> pair = new MutablePair<String, List<Pair<String, Integer>>>();
                pair.setRight(new ArrayList<Pair<String, Integer>>());
                for (Pair<String, List<Pair<String, Integer>>> item : v1) {
                    pair.setLeft(item.getLeft());
                    pair.getRight().addAll(item.getRight());
                }
                for (Pair<String, List<Pair<String, Integer>>> item : v2) {
                    pair.getRight().addAll(item.getRight());
                }
                return Arrays.asList((Pair<String, List<Pair<String, Integer>>>) pair);
            }
        }).collect();

        System.out.println(result.toString());

    }
}
